from utils.postgis import *
from utils.database_table import RemoteSensingRoadTable, RoadTable
import os
from io import StringIO
from pyproj import CRS


def rename_column(df):
    columns = df.columns
    rename_dict = {}
    for col in columns:
        rename_dict[col] = col.lower()
    df.rename(columns=rename_dict, inplace=True)
    return df


def split_multi_linestring(gdf):
    geom_col = gdf.geometry.name
    df_m = gdf[gdf[geom_col].type != "LineString"]
    crs = gdf.crs
    if df_m.shape[0] == 0:
        return gdf
    df_new = []
    for idx in df_m.index:
        geoms = list(gdf.loc[idx, geom_col])
        temp_df = geopd.GeoDataFrame({geom_col: geoms}, geometry=geom_col)
        for col in gdf.columns:  # 复制其他字段
            if col != geom_col:
                temp_df[col] = gdf.loc[idx, col]
        df_new.append(temp_df)
    # 删除multilinestring行
    gdf = gdf.drop(df_m.index)
    df_new.append(gdf)
    gdf = pd.concat(df_new)
    gdf = gdf.set_crs(crs)
    gdf.reset_index(inplace=True, drop=True)
    return gdf


def import_shp(table_name, shp_path):
    if shp_path[-3:] != "shp":
        raise Exception("请输入shp文件路径")
    table = RoadTable(table_name)
    gdf = geopd.read_file(shp_path)
    if gdf.crs is None:
        raise Exception("投影为空")
    gdf.rename_geometry(table.geom, inplace=True)

    # 字段大写转小写
    gdf = rename_column(gdf)

    # 将multilinestring转换为linestring
    if gdf[gdf[table.geom].type != "LineString"].shape[0] != 0:
        gdf = split_multi_linestring(gdf)
        print("完成multilinestring转换为linestring")

    # 投影
    init_crs = gdf.crs
    gdf.to_crs(epsg=SRID, inplace=True)
    print(f"完成投影：{init_crs}->{SRID}")

    # 增加长度列
    gdf[table.length] = gdf[table.geom].length

    gdf.to_postgis(table_name, engine, if_exists="replace")  # 自动创建空间索引
    print(f"成功导入{shp_path}")

    with conn:
        with conn.cursor() as cursor:
            sql = f"""
                alter table {table.name} add column if not exists {table.pk} {table.dtype_dict[table.pk]};
                alter table {table.name} add primary key ({table.pk})"""
            cursor.execute(sql)

            if gdf.shape[0] > 0 and gdf[gdf.geometry.has_z].shape[0] != 0:
                sql = f"""ALTER TABLE {table.name} ALTER COLUMN {table.geom} TYPE geometry(LineString, {SRID}) USING ST_Force2D(geom); """
                cursor.execute(sql)

    table.create_index(table.pk, False)
    print(f"完成创建主键索引")


def load_report_road(shp_dir, prefix):
    # 上报路网
    for filename in os.listdir(shp_dir):
        shp_path = os.path.join(shp_dir, filename)
        table_name, suffix = os.path.splitext(filename)
        if suffix == ".shp":
            table_name = prefix + "_" + table_name[-1:].lower()
            import_shp(table_name, shp_path)


def load_remote_sensing_road(shp_path, table_name):
    import_shp(table_name, shp_path)
